import logging
import unittest

import cv2
import numpy as np
import torch
from PIL import Image
from matplotlib import pyplot as plt
from torch import Tensor
import torch.nn.functional as F

from point_hint_generate.img_util import get_points_image, get_mask_image
from point_hint_generate.loss import dist_loss, balance_loss, label_accuracy_loss
from point_hint_generate.util import predict_mask


class MyTestCase(unittest.TestCase):
    def test_something(self):
        # 固定的输入张量
        input = Tensor([
            [[1.0, 2.0, 3.0], [2.0, 3.0, 4.0], [3.0, 4.0, 5.0], [4.0, 5.0, 6.0], [5.0, 6.0, 7.0]],
            [[-1.0, -2.0, -3.0], [-2.0, -3.0, -4.0], [-3.0, -4.0, -5.0], [-4.0, -5.0, -6.0], [-5.0, -6.0, -7.0]]
        ])
        loss = dist_loss(input)
        print("Loss:", loss.item())

        # 添加断言来验证结果是否符合预期
        # 例如，假设预期的损失是某个确定的值
        expected_loss = 1.4142
        self.assertAlmostEqual(loss.item(), expected_loss, places=4)

    def test_label_accuracy_loss(self):
        # 设置batch_size, points_num, height, width
        batch_size, points_num, height, width = 2, 3, 2, 2

        # 创建input tensor，shape为(batch_size, points_num, 3)
        # 其中前两个维度为归一化的坐标，最后一个维度为置信度
        input = torch.tensor([
            [[0, 0, 0.8], [1, 1, 0.2], [1, 0, 0.9]],  # 第一个batch中的三个点
            [[0, 0, 0.6], [1, 1, 0.1], [0, 1, 0.4]]  # 第二个batch中的三个点
        ])

        # 创建mask tensor，shape为(batch_size, height, width)
        mask = torch.ones(2, 2, 2)

        # 计算损失
        loss = label_accuracy_loss(input, mask)
        print("Loss:", loss.item())

        # 手动计算期望损失值
        expected_loss = F.binary_cross_entropy(torch.tensor([[0.8, 0.2, 0.9], [0.6, 0.1, 0.4]]),
                                               torch.tensor([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]))

        # 验证损失值
        self.assertAlmostEqual(loss.item(), expected_loss.item(), places=4)

    def test_get_img_points(self):
        # Load image using PIL

        image_path = r"C:\Users\crxc\Pictures\car1.png"
        image = cv2.imread(image_path)
        image = torch.tensor(cv2.cvtColor(image, cv2.COLOR_BGR2RGB).transpose(2, 0, 1))

        # image = np.array(image)

        # 创建模拟的dist_loss输入数据
        dummy_coords = np.array([
            [100, 100, 1],  # 正标签
            [120, 130, 0],  # 负标签
            [220, 100, 1],  # 正标签
            [150, 200, 0]  # 负标签
        ])

        dummy_tensor = torch.tensor(dummy_coords)

        # 使用get_points_image方法
        result_image = get_points_image(image, dummy_tensor)

        # 检查结果
        assert np.any(result_image[:, :, :3] != 255)  # 图像应该有其他颜色的像素（由于添加的点）

        # 如果想进一步验证，可以显示返回的图像
        plt.axis('off')
        plt.imshow(result_image)
        plt.show()

    def test_get_mask_image(self):
        image_path = r"C:\Users\crxc\Pictures\car1.png"
        image = cv2.imread(image_path)
        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = torch.tensor(cv2.cvtColor(image, cv2.COLOR_BGR2RGB).transpose(2, 0, 1))
        # 创建模拟的dist_loss输入数据
        dummy_coords = np.array([
            [0.4, 0.4, 1],  # 正标签
            [0.5, 0.5, 1],  # 正标签
            [0.6, 0.6, 1],  # 正标签
            [0.9, 0.9, 0]  # 负标签
        ])
        dummy_tensor = torch.tensor(dummy_coords)
        mask, score = predict_mask(image, dummy_tensor)
        result_image = get_mask_image(image, mask, dummy_tensor)
        plt.axis('off')
        plt.imshow(result_image)
        plt.show()

    def test_log(self):
        logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
        logging.info(f'Using device ')

    def test_dist_loss(self):
        # 定义一个具有确定值的输入张量，用于测试
        batch_size, points_num, coords_dim = 2, 5, 3  # 可以调整这些值来改变测试的批次大小和点的数量
        # 创建一个具有确定坐标的测试输入
        input_tensor = torch.tensor([
            [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0], [3.0, 0.0, 0.0], [4.0, 0.0, 0.0]],
            [[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 2.0, 0.0], [0.0, 3.0, 0.0], [0.0, 4.0, 0.0]]
        ])

        # 调用dist_loss函数
        loss = dist_loss(input_tensor)

        # 检查损失是否为标量
        assert loss.dim() == 0, "Loss is not a scalar."

        # 检查损失是否大于零
        assert loss.item() > 0, "Loss must be positive."

        print(f"Loss: {loss.item()}")
        print("test_dist_loss passed.")

if __name__ == '__main__':
    unittest.main()
